# 导入必要的库
import akshare as ak
import pandas as pd
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
import os
import matplotlib.pyplot as plt
from crewai import Agent, Crew, Task, LLM
import datetime
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="pandas")

# 配置 DeepSeek API 密钥
os.environ["DEEPSEEK_API_KEY"] = "sk-a2b633c6da8c445bbe7ec80b404552b8"  # 替换为实际的 DeepSeek API 密钥

# 初始化 LLM
llm = LLM(
    model="deepseek/deepseek-chat",
    timeout=500,
    api_base="https://api.deepseek.com/v1",
    headers={
        "Authorization": f"Bearer {os.environ['DEEPSEEK_API_KEY']}"
    }
)

# 添加 deepseek-chat 到支持的模型列表
OpenAI._supported_chat_models.append('deepseek-chat')

# 初始化 PandasAI LLM
pandasai_llm = OpenAI(
    model="deepseek-chat",
    api_token=os.environ["DEEPSEEK_API_KEY"],
    api_base="https://api.deepseek.com/v1"
)

# 设置 matplotlib 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 定义函数获取股票数据
# 定义函数获取股票数据
import time
from retrying import retry

# 配置重试条件
def retry_if_connection_error(exception):
    return isinstance(exception, (ConnectionError, TimeoutError))

@retry(retry_on_exception=retry_if_connection_error, wait_fixed=2000, stop_max_attempt_number=3)
def get_stock_data(stock_code, start_date='2024-01-01', end_date='2025-01-01'):
    # 转换日期格式
    start_date = datetime.datetime.strptime(start_date, '%Y-%m-%d')
    end_date = datetime.datetime.strptime(end_date, '%Y-%m-%d')
    
    # 获取股票数据
    try:
        if stock_code.startswith(('sh', 'sz')):
            # A 股数据
            stock_data = ak.stock_zh_a_hist(
                symbol=stock_code[2:],
                start_date=start_date.strftime('%Y%m%d'),
                end_date=end_date.strftime('%Y%m%d'),
                adjust="qfq"
            )
        else:
            # 美股数据
            stock_data = ak.stock_us_daily(symbol=stock_code, start_date=start_date.strftime('%Y-%m-%d'), end_date=end_date.strftime('%Y-%m-%d'), adjust="qfq")
        
        if stock_data.empty:
            print(f"获取 {stock_code} 股票数据为空，请检查代码或股票代码是否正确。")
    except Exception as e:
        print(f"获取 {stock_code} 股票数据失败: {e}")
        raise
    
    # 保存数据为 CSV 文件到桌面
    desktop_path = os.path.join(os.path.expanduser('~'), 'Desktop')
    csv_path = os.path.join(desktop_path, f'{stock_code}_stock_data.csv')
    stock_data.to_csv(csv_path, index=False)
    
    return stock_data

# 定义函数获取相关新闻
def get_stock_news(stock_code):
    try:
        # 获取 A 股新闻
        if stock_code.startswith(('sh', 'sz')):
            stock_name = ak.stock_info_sh_name_code(symbol='A股列表')
            stock_name = stock_name[stock_name['证券代码'] == stock_code[2:]]['证券简称'].values[0]
            news_data = ak.stock_news_em(subject=stock_name)
            if not news_data.empty:
                news_data = news_data['内容'].to_csv(sep='\t', na_rep='nan')
            else:
                news_data = "未获取到相关新闻"
        # 获取美股新闻
        else:
            # 这里可根据实际 API 支持情况修改
            news_data = "美股新闻获取功能暂未完善"
        return news_data
    except Exception as e:
        print(f"获取 {stock_code} 新闻数据失败: {e}")
        return "获取新闻数据失败"

# 定义智能体
analysis_agent = Agent(
    role="股价分析专家",
    goal="根据最近的新闻和股票价格进行股价变动的分析和预测，包括波动率、最大回撤、相关性、趋势判断、收益率计算和异常波动点分析",
    backstory="一位精通金融数据分析的专家，擅长从多维度分析股价变动",
    llm=llm,
    verbose=True
)

chart_agent = Agent(
    role="图表绘制专家",
    goal="绘制股票价格趋势折线图，标注历史最高点和最低点",
    backstory="一位擅长数据可视化的专家，能清晰展示股票价格走势",
    llm=llm,
    verbose=True
)

cause_agent = Agent(
    role="波动成因分析专家",
    goal="分析股票波动的成因",
    backstory="一位深入研究金融市场的专家，能准确剖析股价波动原因",
    llm=llm,
    verbose=True
)

risk_agent = Agent(
    role="风险机会分析专家",
    goal="分析股价的潜在风险和市场机会，为投资决策提供数据支持",
    backstory="一位经验丰富的投资顾问，擅长评估风险和机会",
    llm=llm,
    verbose=True
)

emotion_agent = Agent(
    role="情绪传染分析专家",
    goal="研究情绪传染对股价联动的影响，给出合理推测",
    backstory="一位研究金融市场情绪的专家，能分析情绪对股价的影响",
    llm=llm,
    verbose=True
)

regression_agent = Agent(
    role="回归分析专家",
    goal="构建回归模型分析特定因素对股票价格的影响，给出合理推测并生成可视化图表",
    backstory="一位精通统计分析的专家，擅长构建回归模型",
    llm=llm,
    verbose=True
)

suggestion_agent = Agent(
    role="投资建议专家",
    goal="根据之前智能体的内容总结投资建议",
    backstory="一位资深的投资专家，能综合各方分析给出建议",
    llm=llm,
    verbose=True
)

summary_agent = Agent(
    role="内容总结专家",
    goal="总结以上智能体的内容，以 Markdown 的形式输出结果",
    backstory="一位擅长文档整理和格式化的专家，能将分析结果整合为规范文档",
    llm=llm,
    verbose=True
)

# 主函数
def main():
    # 获取用户输入
    stock_code = input("请输入需要爬取的股票代码（如 sh600000或sz000001）：")

    
    # 获取股票数据
    stock_data = get_stock_data(stock_code)
    
    # 获取相关新闻
    news_data = get_stock_news(stock_code)
    
    # 使用 PandasAI 处理数据
    sdf = SmartDataframe(stock_data, config={"llm": pandasai_llm, "enable_code_execution": True})
    
    # 定义任务
    analysis_task = Task(
        description=f"请根据以下股票数据和新闻进行股价分析：股票数据为 {stock_data.to_csv(sep='	', na_rep='nan')}，新闻数据为 {news_data}。分析股价波动率、最大回撤、与市场指数相关性，用均线判断趋势，计算各周期收益率并找出异常波动点，用中文生成分析结果。",
        agent=analysis_agent,
        expected_output="股价分析结果"
    )

    chart_task = Task(
        description=f"请根据以下股票数据绘制价格趋势折线图：{stock_data.to_csv(sep='	', na_rep='nan')}。以 '日期' 列为横轴，股票价格列作为纵轴。为图表添加标题 '股票价格走势'，横轴标签设为 '日期'，纵轴标签设为 '股票价格'，并添加网格线。同时，在图中标注出股票价格的历史最高点和最低点，标注内容需包含具体价格和对应的日期，用红色标记最高点，绿色标记最低点。",
        agent=chart_agent,
        expected_output="股票价格趋势折线图"
    )

    cause_task = Task(
        description=f"请根据以下股票数据和新闻分析股票波动的成因：股票数据为 {stock_data.to_csv(sep='	', na_rep='nan')}，新闻数据为 {news_data}。用中文生成分析结果。",
        agent=cause_agent,
        expected_output="波动成因分析结果"
    )

    risk_task = Task(
        description=f"请根据以下股票数据和新闻分析股价的潜在风险和市场机会，为投资决策提供数据支持：股票数据为 {stock_data.to_csv(sep='	', na_rep='nan')}，新闻数据为 {news_data}。用中文生成分析结果。",
        agent=risk_agent,
        expected_output="风险和机会分析结果"
    )

    emotion_task = Task(
        description=f"请基于当前股票数据和新闻，不要引入外部数据，研究情绪传染对股价联动的影响，给出合理推测：股票数据为 {stock_data.to_csv(sep='	', na_rep='nan')}，新闻数据为 {news_data}。",
        agent=emotion_agent,
        expected_output="情绪传染分析结果"
    )

    regression_task = Task(
        description=f"请基于当前股票数据，不要引入外部数据，构建回归模型分析特定因素对股票价格的影响，给出合理推测并生成可视化图表：{stock_data.to_csv(sep='	', na_rep='nan')}。",
        agent=regression_agent,
        expected_output="回归分析结果和可视化图表"
    )

    suggestion_task = Task(
        description="请根据之前智能体的输出内容总结投资建议。",
        agent=suggestion_agent,
        context=[analysis_task, chart_task, cause_task, risk_task, emotion_task, regression_task],
        expected_output="投资建议"
    )

    summary_task = Task(
        description="请总结以上智能体的输出内容，以 Markdown 的形式输出结果。",
        agent=summary_agent,
        context=[analysis_task, chart_task, cause_task, risk_task, emotion_task, regression_task, suggestion_task],
        expected_output="Markdown 格式的总结结果"
    )

    # 创建并运行工作流
    workflow_crew = Crew(
        agents=[analysis_agent, chart_agent, cause_agent, risk_agent, emotion_agent, regression_agent, suggestion_agent, summary_agent],
        tasks=[analysis_task, chart_task, cause_task, risk_task, emotion_task, regression_task, suggestion_task, summary_task],
        verbose=True
    )

    result = workflow_crew.kickoff()

    # 输出结果
    if hasattr(result, 'content'):
        result_content = result.content
    else:
        result_content = str(result)

    print(result_content)

    # 保存结果为 Markdown 文件
    with open("stock_analysis_summary.md", "w", encoding="utf-8") as f:
        f.write(result_content)

if __name__ == "__main__":
    main()